library(plyr)
library(estimatr)
library(emmeans)
library(tidyverse)
library(magrittr)
library(broom)
library(specr)

d1 <- "https://github.com/thomasjwood/full_fact/raw/main/ff_mt.rds" %>% 
  url %>% 
  gzcon %>% 
  readRDS %>% 
  mutate(
    out_agree = out_agree %>% 
      mapvalues(
        c("Tend to disagree",
          "Tend to agree"),
        c("Disagree",
          "Agree")
      ) %>% 
      factor(
        c("Strongly disagree",
          "Disagree",
          "Neither agree nor disagree",
          "Agree",
          "Strongly agree")
      ),
    out_true = out_true %>% 
      factor(
        c("True",
          "Probably true",
          "Not sure",
          "Probably false",
          "False") %>% 
          rev
      ),
    agree_num = out_agree %>% 
      as.numeric,
    true_num = out_true %>% 
      as.numeric,
    cond = cond %>% 
      factor(
        c("cond_items",
          "cond_misinfo",
          "cond_corr")
      )
  ) %>% 
  filter(
    issue %>% 
      str_detect(
        "saltwater|cooling"
      ) &
      wave == "wave 1"
  ) %>% 
  mutate(
    com_num = agree_num %>% 
      add(
        true_num
      ) %>% 
      divide_by(2)
  )

t1 <- d1 %>% 
  group_by(
    country, issue
  ) %>% 
  nest

t1$mods <- t1$data %>% 
  map(
    ~lm_robust(
      com_num ~ cond, 
      data = .x
    )
  )

t1$emm <- t1$mods %>% 
  map2(
    t1$data,
    ~emmeans(
      .x, 
      consec ~ cond, 
      data = .y
    )
  )

t1$edf <- t1$emm %>% 
  map(
    ~extract2(., 1) %>% 
      tidy
  )

t1$cont <- t1$emm %>% 
  map(
    ~extract2(., 2) %>% 
      tidy
  )

# first, fitted values

t2 <- t1 %>% 
  ungroup %>% 
  select(
    country, issue, edf 
  ) %>% 
  pmap_dfr(
    function(
      country, issue, edf
    )
      
      edf %>% 
      mutate(
        issue = issue, 
        country = country
      )
  ) %>% 
  mutate(
    horiz = "Conditional means" 
  ) %>% 
  bind_rows(
    t1 %>% 
      ungroup %>% 
      select(
        country, issue, cont
      ) %>% 
      pmap_dfr(
        function(
          country, issue, cont
        )
          
          cont %>% 
          mutate(
            issue = issue, 
            country = country
          )
      ) %>% 
      rename(p.value = "adj.p.value") %>% 
      mutate(
        horiz = contrast  %>% 
          str_detect("items") %>% 
          ifelse(
            "Misinformation effect",
            "Correction effect"
          )
      ) %>% 
      select(-term, -contrast, -null.value)
  ) %>% 
  mutate(
    lo = estimate %>%
      subtract(
        std.error %>%
          multiply_by(2.58)
      ),
    hi = estimate %>%
      add(
        std.error %>%
          multiply_by(2.58)
      ),
    lab = estimate %>%
      round(2) 
  )

t2$lab <- t2$cond %>% 
  is.na %>% 
  not %>% 
  ifelse(
    t2$lab,
    t2$lab %>% 
      str_c(
        gtools::stars.pval(
          t2$p.value
        )
      )
  )

t2$lab %<>% str_trim 

t2$lab %<>% 
  equals("0") %>% 
  ifelse(
    t2$estimate %>% 
      round(4) %>% 
      prettyNum(scientific = F),
    t2$lab
  )

# remove leading zeroes

t2$lab %<>% 
  str_extract("\\d") %>% 
  equals("0") %>% 
  ifelse(
    t2$lab %>% 
      str_replace("0\\.", "\\."),
    t2$lab
  )

# add decimal to missing decimals

t2$lab %<>% 
  str_detect(
    "\\."
  ) %>% 
  ifelse(
    t2$lab,
    t2$lab %>% 
      str_replace(
        '\\d', 
        t2$lab %>% 
          str_extract("\\d") %>% 
          str_c("\\.")
      )
  )

# add trailing zero

t2$lab  %<>%  
  str_count("\\d") %>% 
  equals(1) %>% 
  ifelse(
    t2$lab %>% 
      str_replace(
        "\\d",
        t2$lab %>% 
          str_extract("\\d") %>% 
          str_pad(width = 2, side = "right", pad = "0")
      ),
    t2$lab
  )

t2$country %<>% 
  factor(
    t2 %>% 
      filter(
        cond %>% 
          is.na %>% 
          not
        ) %>% 
      group_by(country) %>% 
      summarize(
        e2 = estimate %>% mean
        ) %>% 
      arrange(
        e2
      ) %>% 
      use_series(country)
    )

t2$horiz %<>% 
  factor(
    c("Conditional means",
      "Misinformation effect",
      "Correction effect")
  )

t2$cy <- case_when(
  t2$cond %>% 
    is.na %>% 
    not &
  t2$cond %>% 
    equals("cond_misinfo") ~
    t2$country %>% 
    as.numeric %>% 
    add(.15),
  t2$cond %>% 
    is.na %>% 
    not &
    t2$cond %>% 
    equals("cond_corr") ~
    t2$country %>% 
    as.numeric %>% 
    subtract(.15),
  t2$cond %>% 
    is.na %>% 
    not &
    t2$cond %>% 
    equals("cond_items") ~
    t2$country %>% 
    as.numeric,
  TRUE ~ NA_real_
  )

t2$issue %<>%
  mapvalues(
    c("global_cooling", 
      "saltwater_covid"),
      c("Global cooling '16-18",
      "Saltwater kills COVID")
    )

t2$lab[
  t2$cond %>% is.na %>% not
  ] %<>% 
  str_replace(
    "\\.\\d{2}",
    t2$lab[
      t2$cond %>% is.na %>% not
    ] %>% 
      str_extract(
        "\\.\\d{1}"
      )
  )

t2$country %<>% 
  mapvalues("nigeria", "Nigeria")

t2 %>% 
  filter(
    horiz %>% 
      str_detect("Conditional ", negate = T)
  ) %>% 
  ggplot() +
  geom_blank(
    aes(y = country, color = cond,
        xmin = lo, xmax = hi),
    data = t2 %>% 
      filter(
        cond %>% 
          is.na %>% 
          not &
        horiz %>% 
          str_detect("Conditional ", negate = T)
        ),
    position = position_dodge(
      width = .85
    ),
    show.legend = F
  ) +
  geom_point(
    aes(y = country, 
        x = estimate,
        color = cond,
    ),
    data = t2 %>% 
      filter(
        horiz %>% 
          str_detect("Conditional ", negate = T) &
        cond %>% 
          is.na %>% 
          not
      ),
    position = position_dodge(
      width = .45
    ),
    size = 6.4,
    shape = 21) +
  geom_vline(
    aes(
      xintercept = xint
    ),
    t2 %>% 
      filter(
        cond %>% 
          is.na &
        horiz %>% 
          str_detect("Conditional ", negate = T)
        
      ) %>% 
      mutate(xint = 0),
    size = .25,
    linetype = "dashed"
  ) +
  geom_linerange(
    aes(
      y = country, xmin = lo, xmax = hi
    ),
    t2 %>% 
      filter(
        cond %>% 
          is.na &
        horiz %>% 
          str_detect("Conditional ", negate = T)
        )
  ) +
  geom_point(
    aes(
      estimate,
      country, 
      group = cond
    ),
    shape = 21,
    size = 8,
    fill = "white",
    t2 %>% 
      filter(
        cond %>% 
          is.na &
        horiz %>% 
          str_detect("Conditional ", negate = T)
        ),
    position = position_dodge(
      width = .75
    ),
    
  ) +
  geom_text(
    aes(
      estimate,
      country, 
      label = lab %>% 
        str_replace(fixed("0."), ".")
    ),
    data = t2 %>% 
      filter(
        cond %>% 
          is.na &
        horiz %>% 
          str_detect("Conditional ", negate = T)
        ),
    size = 2
    ) +
  facet_grid(
    issue ~ horiz, 
    scales = "free",
    space ="free",
    labeller = label_wrap_gen(width = 10)
  ) +
  scale_color_grey() +
  scale_x_continuous() +
  scale_y_discrete(
    breaks = t2$country %>% 
      levels,
    labels = t2$country %>% 
      levels %>% 
      str_replace("\\s", "\n"),
  ) +
  labs(
    x = "",
    y = "",
    color = ""
  ) +
  theme_minimal(
    base_family = "Roboto" 
  ) +
  theme(
    strip.text.x = element_text(
      # size =  11, 
      angle = 0
    ),
    strip.text.y = element_text(
      # size =  11, 
      angle = 0
    ),
    legend.position = "none",
    axis.text.y = element_text(
      # angle = 45/2,
      hjust = .5
      # vjust = 1
    ),
    plot.margin = margin(.2, .2, .2 , 1, "cm") 
  )
